"""
创建一个组件，它的作用是 重写 用户的提问

核心的方法是 rewrite
"""

import re
import json
from typing import List, Dict, Any, Optional
from nl2sql.engines.parse_query import ParseQuery
from nl2sql.model.llm import LLM
from nl2sql.datasource.excel_datasource import ExcelDatasource
from nl2sql.model.embedding_model import EmbeddingModel
from nl2sql.prompt.query_rewrite_prompt import query_rewrite_prompt


class QueryRewriter:
    def __init__(self,
                 embedding_model: EmbeddingModel,
                 llm: LLM,
                 data_source: ExcelDatasource ):
        """
        __init__() 这个函数叫 构造函数（初始化函数）
        作用就是创建、定义这个函数的基本属性
        """

        self.llm: LLM = llm  # 声明 self.llm 是一个 LLM 类
        self.embedding_model: EmbeddingModel = embedding_model  # 声明self.embedding_model 这个变量，是一个EmbeddingModel

        self.data_source: ExcelDatasource = data_source

        if self.embedding_model:
            self.data_source.embedding_model = self.embedding_model

        self.parse_query = ParseQuery(model=self.llm)

        self.rewrite_prompt = query_rewrite_prompt

    def parse_and_find_closest_enums(self, query: str) -> List[Dict]:
        """
        解析query并提取最接近的枚举值列表
        """
        parsed: str = self.parse_query.parse_query(query=query)

        # 添加了容错机制，提升整体的稳定性
        try:
            parsed: dict = json.loads(parsed)
        except json.JSONDecodeError:
            return []

        if not parsed or "values" not in parsed:
            raise ValueError("Parser结果无效，缺少 'values' 字段")

        simi_info = []
        for value in parsed["values"]:
            closest_enums = {}
            result = self.data_source.find_closest_enum(input_value=value)
            if result and 'closest_enum' in result:
                closest_enums['原始内容'] = value
                closest_enums['可能的信息'] = result['closest_enum']
                simi_info.append(closest_enums)

        return simi_info

    @staticmethod
    def postprocess(text: str) -> str:
        """
        后处理大模型输出，只提取第一个有效字符串
        """
        text = text.strip()
        match = re.search(r"[\u4e00-\u9fa5a-zA-Z0-9].*", text)
        if match:
            return match.group(0).strip()
        else:
            return text

    def rewrite_query(self, query: str) -> str:
        """
        完整执行一轮query改写流程
        """
        closest_enums_infos = self.parse_and_find_closest_enums(query=query)

        print(closest_enums_infos)
        rewrite_prompt = self.rewrite_prompt.format(query=query,
                                                    closest_enums=closest_enums_infos)

        rewritten_query = self.llm.chat(prompt=rewrite_prompt)

        return rewritten_query


if __name__ == "__main__":
    query = "载客输出近一个月的华航，东航，大韩航空，海航这四家中，从江城，天府起飞的航班载客量"

    apikey = "sk-68ac5f5ccf3540ba834deeeaecb48987"
    emb_model = EmbeddingModel(api_key=apikey)

    api_key: str = 'sk-68ac5f5ccf3540ba834deeeaecb48987'
    base_url: str = "https://dashscope.aliyuncs.com/compatible-mode/v1"
    model_name: str = "deepseek-v3"

    llm = LLM(api_key=api_key,
              base_url=base_url,
              model_name=model_name,
              temperature=0)

    data_path = 'D:/PythonProject_airline/nl2sql/examples/data'
    excel_ds = ExcelDatasource(source_path=data_path, embedding_model=emb_model)

    query_rewriter = QueryRewriter(embedding_model=emb_model,
                                   llm=llm,
                                   data_source=excel_ds)

    rewritten_query = query_rewriter.rewrite_query(query=query)

    print("改写后的新Query：")
    print(rewritten_query)
